from reportlab.lib import colors, pagesizes
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.platypus import SimpleDocTemplate, Paragraph, Table, TableStyle, PageBreak

from ai_translator_v2.book import Book, ContentType
from ai_translator_v2.utils import LOG


class Writer(object):
    """
    用于保存翻译后书籍的类。
    """

    def __init__(self):
        """
        初始化Writer类。
        """
        pass

    def save_translated_book(self, book: Book, output_file_format: str, output_file_path: str = None, ):
        """
        根据指定的输出格式保存翻译后的书籍。

        :param output_file_path:
        :param book: 被翻译的书籍对象。
        :param output_file_format: 输出文件格式，支持"pdf"和"markdown"。
        :return: 保存翻译后书籍的文件路径。
        """
        LOG.debug(f"output_file_format: {output_file_format}")
        if output_file_format.lower() == "pdf":
            output_file_path = self._save_translated_book_pdf(book, output_file_path)
        elif output_file_format.lower() == "markdown":
            output_file_path = self._save_translated_book_markdown(book, output_file_path)
        else:
            LOG.error(f"Not supported filetype: {output_file_format}")
        LOG.info(f"Translated completed. file saved to {output_file_path}")
        return output_file_path

    @staticmethod
    def _save_translated_book_pdf(book: Book, output_file_path: str = None):
        """
        将翻译后的书籍保存为PDF格式。

        :param book: 被翻译的书籍对象。
        :return: 保存翻译后书籍的PDF文件路径。
        """
        if output_file_path is None:
            output_file_path = book.pdf_file_path.replace(".pdf", "_translated.pdf")
        LOG.info(f"开始导出: {output_file_path}")

        # 注册中文字体
        font_path = "../fonts/simsun.ttc"  # 字体文件路径
        pdfmetrics.registerFont(TTFont("SimSun", font_path))

        # 创建PDF文档样式
        simsun_style = ParagraphStyle('SimSun', fontName='SimSun', fontSize=12, leading=14)

        # 创建PDF文档
        doc = SimpleDocTemplate(output_file_path, pagesize=pagesizes.letter)
        styles = getSampleStyleSheet()
        story = []

        # 遍历页面和内容，添加到PDF中
        for page in book.pages:
            for content in page.contents:
                if content.status:
                    if content.content_type == ContentType.TEXT:
                        # 添加翻译文本到PDF
                        text = content.translation
                        para = Paragraph(text, simsun_style)
                        story.append(para)

                    elif content.content_type == ContentType.TABLE:
                        # 添加表格到PDF
                        table = content.translation
                        table_style = TableStyle([
                            ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
                            ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
                            ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
                            ('FONTNAME', (0, 0), (-1, 0), 'SimSun'),  # 表头字体设置为 "SimSun"
                            ('FONTSIZE', (0, 0), (-1, 0), 14),
                            ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
                            ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
                            ('FONTNAME', (0, 1), (-1, -1), 'SimSun'),  # 表格字体设置为 "SimSun"
                            ('GRID', (0, 0), (-1, -1), 1, colors.black)
                        ])
                        pdf_table = Table(table.values.tolist())
                        pdf_table.setStyle(table_style)
                        story.append(pdf_table)
            # 在每个页面后添加分页符，除了最后一页
            if page != book.pages[-1]:
                story.append(PageBreak())

        # 生成并保存PDF文件
        doc.build(story)
        return output_file_path

    @staticmethod
    def _save_translated_book_markdown(book: Book, output_file_path: str = None):
        """
        将翻译后的书籍保存为Markdown格式。

        :param book: 被翻译的书籍对象。
        :return: 保存翻译后书籍的Markdown文件路径。
        """
        if output_file_path is None:
            output_file_path = book.pdf_file_path.replace('.pdf', '_translated.md')
        LOG.info(f"开始导出: {output_file_path}")
        with open(output_file_path, 'w', encoding='utf-8') as output_file:
            for page in book.pages:
                for content in page.contents:
                    if content.status:
                        if content.content_type == ContentType.TEXT:
                            # 添加翻译文本到Markdown文件
                            text = content.translation
                            output_file.write(text + '\n\n')

                        elif content.content_type == ContentType.TABLE:
                            table = content.translation
                            header = '| ' + ' | '.join(str(column) for column in table.columns) + ' |' + '\n'
                            separator = '| ' + ' | '.join(['---'] * len(table.columns)) + ' |' + '\n'
                            body = '\n'.join(['| ' + ' | '.join(str(cell) for cell in row) + ' |' for row in
                                              table.values.tolist()]) + '\n\n'
                            output_file.write(header + separator + body)

                # 在每个页面后添加分页符（水平线），除了最后一页
                if page != book.pages[-1]:
                    output_file.write('---\n\n')
        return output_file_path
